"""
connect to robot
"""
import base64, io
import functools

import paramiko
from shared_data import ROBOT_KEY_STR
from typing import Optional
from paramiko.ssh_exception import BadAuthenticationType
from app.core.exceptions import RobotException, ErrorCode, ParamikoException
import select, time, sys
from dataclasses import dataclass, field

USE_LOG = False

if USE_LOG:
    from app.core.config.__logging__ import get_logger

    logger = get_logger("driver.paramiko")
else:
    class FakeLogger:

        @staticmethod
        def info(msg: str):
            print(msg)


    logger = FakeLogger()


@dataclass
class ExecuteCommandInterface:
    status: bool
    output: str
    error: str = field(default="")


def is_connected(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if self.ssh_client is None or self.sftp_client is None or self.shell_client is None:
            raise RobotException(ErrorCode.ROBOT_CONNECT_FAIL, "Connection closed")
        return func(self, *args, **kwargs)

    return wrapper


class ParamikoDriver:
    def __init__(self, user_name: str = "root", port: int = 22, password: str = "None", time_out=10):
        self.host = None
        self.port = port
        self.username = user_name
        self.password = password
        self.time_out = time_out
        self.ssh_client: Optional[paramiko.SSHClient] = None
        self.sftp_client: Optional[paramiko.SFTPClient] = None
        self.shell_client = None

    def connect(self, robot_ip: str) -> bool:
        self.host = robot_ip
        self.ssh_client = paramiko.SSHClient()
        self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            self.ssh_client.connect(self.host, port=self.port, username=self.username,
                                    password=self.password, timeout=self.time_out)
            self.sftp_client = self.ssh_client.open_sftp()
            self.shell_client = self.ssh_client.invoke_shell()
            logger.info("Robot Connected !")
            return True
        except BadAuthenticationType:
            # connet robot by key_str
            key_data = base64.b64decode(ROBOT_KEY_STR.strip())
            key = paramiko.RSAKey.from_private_key(io.StringIO(key_data.decode("utf-8")))
            try:
                self.ssh_client.connect(self.host, port=self.port, username=self.username, password=self.password,
                                        timeout=self.time_out, pkey=key)
                self.sftp_client = self.ssh_client.open_sftp()
                self.shell_client = self.ssh_client.invoke_shell()
                logger.info("Robot Connected !")
                return True
            except BadAuthenticationType as e:
                raise RobotException(error_code=ErrorCode.ROBOT_CONNECT_FAIL,
                                     detail=f"Probably need to upload the public key to this robot, {e}")
        except InterruptedError:
            return False
        except Exception as error:
            raise RobotException(error_code=ErrorCode.ROBOT_CONNECT_FAIL, detail=f"Connect to robot fail, {error}")

    def close(self):
        try:
            if self.ssh_client:
                self.ssh_client.close()
            if self.sftp_client:
                self.sftp_client.close()
            self.ssh_client = None
            self.sftp_client = None
            logger.info("Robot Closed!")
        except Exception as error:
            raise RobotException(error_code=ErrorCode.ROBOT_DISCONNECT_FAIL, detail=str(error))

    @is_connected
    def _read_available(self) -> str:
        """读取所有可用输出"""
        output = ""
        while self.shell_client.recv_ready():
            output += self.shell_client.recv(1024).decode('utf-8')
        return output

    @is_connected
    def execute_shell_command(self, cmd: str, timeout=None):
        """
        run command in a robot, keep reading the output
        :param cmd:
        :param timeout: if timeout is None, ignore the timeout
        :return:
        """
        self.start_cli(send_command=cmd, timeout=timeout)

    @is_connected
    def start_cli(self, send_command: Optional[str] = None, timeout=None):
        def _is_waiting_for_input(output):
            """通过输出内容特征判断是否在等待用户输入"""
            waiting_patterns = [
                "enter.*name",  # 匹配 "enter OPERATOR name" 等
                "password:",
                "please input",
                "input.*:",
                "y/n",  # 确认提示
                "continue.*?",  # 继续提示
                "select.*option",  # 选择选项
                "press.*key",  # 按键提示
                "username:",  # 用户名
                "login:",  # 登录
                ":",  # 冒号提示（通用）
                ">",  # 大于号提示
                "\?$"  # 以问号结尾
                "exit"
                "logout"
            ]
            last_line = output.split('\n')[-1].strip().lower()
            for pattern in waiting_patterns:
                import re
                if re.search(pattern, last_line, re.IGNORECASE):
                    return True
            return False

        logger.info("进入交互模式，输入 'exit' 或 'quit' 退出")
        while True:
            if send_command is None:
                # 显示命令提示符
                sys.stdout.write("$ ")
                sys.stdout.flush()
                # 读取用户输入
                command = input().strip()
                if command.lower() in ['exit', 'quit']:
                    self.shell_client.send('\x03')
                    time.sleep(0.5)  # 等待中断处理
                elif not command:
                    self.shell_client.send("\x0d")
            else:
                command = send_command
            self.shell_client.send(command + "\n")
            start_time = time.time()
            output_buffer = ""
            while True:
                if timeout is None:
                    pass
                else:
                    if time.time() - start_time > timeout:
                        print("命令执行超时")
                        break

                # 使用select检查是否有数据可读
                time.sleep(0.1)
                if select.select([self.shell_client], [], [], 0.1)[0]:
                    try:
                        # 读取数据
                        data = self.shell_client.recv(9096).decode('utf-8')
                        output_buffer += data
                        # 实时打印新数据
                        if data:
                            # 去除回显
                            try:
                                if command.strip() == "" and "# " in data:
                                    data = data.split("\r\n")[-1]
                                elif command.strip() != "":
                                    data = data.replace(f"{command}\r\n", "")
                                else:
                                    pass
                            except Exception:
                                pass
                            print(data, end='', flush=True)
                        if _is_waiting_for_input(output_buffer):
                            break
                        if command.lower() in ['exit', 'quit']:
                            break

                    except Exception as e:
                        print(f"读取错误: {e}")
                        break
            if command.lower() in ['exit', 'quit']:
                break

    @is_connected
    def execute_command(self, command: str, time_out=30) -> ExecuteCommandInterface:
        """
        执行一次会话的命令，
        :param command: 执行的命令，如 date/pwd
        :param time_out:
        :return:
        """
        try:
            logger.info(f"Execute {command}")
            stdin, stdout, stderr = self.ssh_client.exec_command(command)
            exit_status = stdout.channel.recv_exit_status()
            output = stdout.read().decode('utf-8').strip()
            error = stderr.read().decode('utf-8').strip()
            if not exit_status==0:
                raise ParamikoException(ErrorCode.PARAMIKO_EXECUTE_FAIL, detail=f"execute {command} fail \r\n" + error)
            return ExecuteCommandInterface(
                status=exit_status==0,
                output=output,
                error=error
            )
        except RobotException:
            raise
        except Exception as e:
            raise ParamikoException(ErrorCode.PARAMIKO_EXECUTE_FAIL, detail=f"execute {command} fail \r\n" + str(e))


if __name__ == '__main__':
    obj = ParamikoDriver(host="192.168.6.95")
    obj.connect()
    result = obj.execute_command("pwd")
